package gossip.v1;

import org.ricks.net.MessageProcessor;
import org.ricks.net.StateMachineEnum;
import org.ricks.net.extension.protocol.StringProtocol;
import org.ricks.net.transport.AioQuickClient;
import org.ricks.net.transport.AioQuickServer;
import org.ricks.net.transport.AioSession;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

public class GossipTransport {
    private final int port;
    private final GossipCore core;
    private AioQuickServer server;
    private final Map<InetSocketAddress, SessionHolder> sessionCache = new ConcurrentHashMap<>();
    private final ExecutorService connectionPool = Executors.newFixedThreadPool(4);

    // 在类级别定义常量
    private static final int MAX_SESSION_QUEUE = 1000;

    public GossipTransport(int port, GossipCore core) {
        this.port = port;
        this.core = core;
    }

    public void start() {
        server = new AioQuickServer(port, new StringProtocol(), new MessageProcessor<String>() {
            @Override
            public void process(AioSession session, String msg) {
                try {
                    core.handleMessage(session.getRemoteAddress(), msg);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

            @Override
            public void stateEvent(AioSession session, StateMachineEnum state, Throwable throwable) {
                if (throwable != null) {
                    try {
                        System.err.printf("[NETWORK] Error in session %s: %s%n",
                                session.getRemoteAddress(), throwable.getMessage());
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                }
            }
        });

        try {
            server.setReadBufferSize(1024)
                    .setThreadNum(2)
                    .start();
            System.out.printf("[NETWORK] Transport started on port %d%n", port);
        } catch (IOException e) {
            throw new RuntimeException("Failed to start transport", e);
        }
    }

    public void send(InetSocketAddress target, String data) {
        connectionPool.execute(() -> {
            try {
                SessionHolder holder = sessionCache.computeIfAbsent(target, k ->
                        new SessionHolder(target));

                holder.send(data);
            } catch (Exception e) {
                System.err.printf("[NETWORK] Failed to send to %s: %s%n", target, e.getMessage());
                sessionCache.remove(target);
            }
        });
    }

    // 内部类访问外部类的成员变量
    private class SessionHolder {
        private final InetSocketAddress address;
        private AioSession session;
        private final Queue<String> pendingMessages = new ConcurrentLinkedQueue<>();
        private final AtomicInteger pendingCount = new AtomicInteger(0);
        private volatile boolean isConnecting;

        public SessionHolder(InetSocketAddress address) {
            this.address = address;
        }

        public synchronized void send(String data) {
            if (session != null && session.isInvalid()) {
                session = null;
            }

            if (session != null) {
                try {
                    // 使用正确的write方法
                    session.writeBuffer().writeInt(4 + data.getBytes().length);
                    session.writeBuffer().write(data.getBytes());
                    session.writeBuffer().flush();
                    return;
                } catch (Exception e) {
                    session = null;
                }
            }

            pendingMessages.offer(data);
            if (pendingMessages.size() > MAX_SESSION_QUEUE) {
                pendingMessages.poll(); // Drop oldest message to prevent OOM
            }

            if (!isConnecting) {
                isConnecting = true;
                // 使用外部类的connectionPool
                GossipTransport.this.connectionPool.execute(this::connectAndSend);
            }
        }

        private void connectAndSend() {
            try {
                // 使用外部类的core引用
                AioQuickClient client = new AioQuickClient(address.getHostName(),
                        address.getPort(), new StringProtocol(), new MessageProcessor<String>() {
                    @Override
                    public void process(AioSession session, String msg) {
                        // 处理响应
                        try {
                            GossipTransport.this.core.handleMessage(session.getRemoteAddress(), msg);
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    }

                    @Override
                    public void stateEvent(AioSession session, StateMachineEnum state, Throwable throwable) {
                        if (throwable != null) {
                            System.err.printf("[NETWORK] Client error: %s%n", throwable.getMessage());
                        }
                    }
                });
                session = client.start();

                // 发送所有待处理消息
                String msg;
                while ((msg = pendingMessages.poll()) != null) {
                    // 使用正确的write方法
                    session.writeBuffer().writeInt(4 + msg.getBytes().length);
                    session.writeBuffer().write(msg.getBytes());
                    session.writeBuffer().flush();
                }
            } catch (Exception e) {
                System.err.printf("[NETWORK] Connection to %s failed: %s%n", address, e.getMessage());
                session = null;
            } finally {
                isConnecting = false;

                // 如果还有待处理消息，再次尝试连接
                if (!pendingMessages.isEmpty() && !isConnecting) {
                    isConnecting = true;
                    GossipTransport.this.connectionPool.execute(this::connectAndSend);
                }
            }
        }
    }
}
